"""
model name : 目标跟踪
file       : track.py
information:
    author : OuYang
    time   : 2025/7/19
"""
import cv2
import numpy as np
from ultralytics import YOLO


def track_yolo():
    """
    使用Ultralytics自带的追踪器进行目标跟踪
    :return:
    """
    model = YOLO("runs/detect/yolov8n_otb_car24/weights/best.pt")

    result = model.track(
        source="test_data/Car24_val.avi",
        show=True,
    )


def track_kalman():
    """
    使用cv2的卡尔曼滤波跟踪
    :return:
    """

    # 初始化卡尔曼滤波
    kalman_xy = cv2.KalmanFilter(4, 2)
    kalman_xy.measurementMatrix = np.array([[1, 0, 0, 0], [0, 1, 0, 0]], np.float32)
    # 观测矩阵
    kalman_xy.transitionMatrix = np.array([
        [1, 0, 1, 0],
        [0, 1, 0, 1],
        [0, 0, 1, 0],
        [0, 0, 0, 1]
    ], np.float32)
    kalman_xy.processNoiseCov = 1e-4 * np.eye(4, dtype=np.float32)

    kalman_wh = cv2.KalmanFilter(4, 2)
    kalman_wh.measurementMatrix = np.array([[1, 0, 0, 0], [0, 1, 0, 0]], np.float32)
    kalman_wh.transitionMatrix = np.array([
        [1, 0, 1, 0],
        [0, 1, 0, 1],
        [0, 0, 1, 0],
        [0, 0, 0, 1]
    ], np.float32)
    kalman_wh.processNoiseCov = 1e-4 * np.eye(4, dtype=np.float32)

    # 打开目标检测
    model = YOLO("runs/detect/yolov8n_otb_car24/weights/best.pt")

    # 打开视频
    capture = cv2.VideoCapture("test_data/Car24_val.avi")
    if not capture.isOpened():
        print("Error opening video stream or file")
        exit(-1)

    _, frame = capture.read()
    video_width = frame.shape[1]
    video_height = frame.shape[0]

    # 绘制第一帧中检测目标
    from utils.opencv import RectangleDrawer
    drawer = RectangleDrawer(frame)
    drawer.run()
    first_post = drawer.get_rectangle_info()
    kalman_xy.statePost = np.array([
        [first_post["center"][0]],
        [first_post["center"][1]],
        [0], [0]], np.float32
    )
    first_frame_box_width = first_post["width"]
    first_frame_box_height = first_post["height"]
    kalman_wh.statePost = np.array([
        [first_frame_box_width],
        [first_frame_box_height],
        [0], [0]], np.float32
    )

    predictions = []
    measurements = []

    fi = 0
    # 读取视频
    while True:
        fi += 1

        ret, frame = capture.read()
        if not ret:
            break

        # 在100帧和125帧之间模拟目标消失
        if 100 < fi < 125:
            cv2.rectangle(frame, (0, 0), (video_width, video_height), (0, 0, 255), -1)
        # 在200帧和250帧之间模拟目标消失
        if 200 < fi < 250:
            cv2.rectangle(frame, (0, 0), (video_width, video_height), (0, 0, 255), -1)

        # 目标检测，获取检测框信息
        results = model(frame)
        if len(results[0]) == 1:
            detect_xywh = [int(i) for i in results[0].boxes[0].xywh[0].tolist()]
        else:
            detect_xywh = [0, 0, 0, 0]
        measurements.append(detect_xywh)

        # 卡尔曼滤波产生预测位置
        prediction = kalman_wh.predict()
        predict_wh = [int(i) for i in prediction[:2].flatten()]
        prediction = kalman_xy.predict()
        predict_xy = [int(i) for i in prediction[:2].flatten()]
        predictions.append(predict_xy)

        # 卡尔曼滤波更新
        kalman_xy.correct(np.array([detect_xywh[0], detect_xywh[1]], np.float32))
        kalman_wh.correct(np.array([detect_xywh[2], detect_xywh[3]], np.float32))
        if len(results[0]) == 1:  # 如果有目标则用，目标检测得到的结果作为卡尔曼滤波更新
            kalman_xy.correct(np.array([detect_xywh[0], detect_xywh[1]], np.float32))
            kalman_wh.correct(np.array([detect_xywh[2], detect_xywh[3]], np.float32))
        else:  # 当目标检测没有找到目标时，使用本次预测值作为下一次更新
            kalman_xy.correct(np.array([[predict_xy[0]], [predict_xy[1]]], np.float32))
            kalman_wh.correct(np.array([[predict_wh[0]], [predict_wh[1]]], np.float32))

        # 绘制矩形框, 目标检测框
        cv2.rectangle(frame,
                      (detect_xywh[0] - detect_xywh[2] // 2, detect_xywh[1] - detect_xywh[3] // 2),
                      (detect_xywh[0] + detect_xywh[2] // 2, detect_xywh[1] + detect_xywh[3] // 2),
                      (255, 0, 0), 1
                      )

        # 绘制矩形框，卡尔曼滤波
        cv2.rectangle(frame,
                      (predict_xy[0] - predict_wh[0] // 2, predict_xy[1] - predict_wh[1] // 2),
                      (predict_xy[0] + predict_wh[0] // 2, predict_xy[1] + predict_wh[1] // 2),
                      (0, 255, 0), 1
                      )
        cv2.imshow('tracking', frame)
        if cv2.waitKey(25) & 0xFF == ord('q'):
            break
    cv2.destroyAllWindows()

    # 绘制, 轨迹
    image = np.zeros((video_height, video_width, 3), np.int8)

    for measurement in measurements:
        cv2.circle(image, (int(measurement[0]), int(measurement[1])), 1, (0, 0, 255), -1)

    for prediction in predictions:
        cv2.circle(image, (int(prediction[0]), int(prediction[1])), 1, (0, 255, 0), -1)

    cv2.imshow("image", image)
    cv2.waitKey(0)
    cv2.destroyAllWindows()


if __name__ == "__main__":
    # track_yolo()
    track_kalman()
